import torch
import numpy as np
import torch.nn.functional as F
from torch.distributions import Categorical


# 创建一个 NumPy 数组
state = np.array([1, 2, 3, 4])

# 转换为 PyTorch 张量
tensor_state = torch.from_numpy(state)

# 转换为浮点型
float_tensor_state = tensor_state.float()

# 在第 0 维上增加一个维度
unsqueezed_tensor_state = float_tensor_state.unsqueeze(0)

print("原始 NumPy 数组:", state)
print("转换后的 PyTorch 张量:", tensor_state)
print("转换为浮点型后的张量的数据类型:", float_tensor_state.dtype)
print("原始张量的形状:", float_tensor_state.shape)
print("增加维度后的张量的形状:", unsqueezed_tensor_state.shape)
print("增加维度后的张量:", unsqueezed_tensor_state)


# 假设 action_logits 是神经网络的输出
action_logits = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])

# 应用 softmax 函数
action_probs = F.softmax(action_logits, dim=-1)

print("原始的 action_logits:")
print(action_logits)
print("应用 softmax 后的 action_probs:")
print(action_probs)

m = Categorical(action_probs)
action = m.sample()
log_prob = m.log_prob(action)

print("采样的动作:", action)
print("动作的对数概率:", log_prob)